import PIL
import matplotlib
from PIL import Image
from captcha.image import ImageCaptcha
import matplotlib.pyplot as plt
import numpy as np
import random
import os

from tf_seq_generator import BaseSequence

os.environ["TF_CPP_MIN_LOG_LEVEL"]='2' # 只显示 warning 和 Error

from tensorflow.python.keras.utils import Sequence
import string
characters = string.digits + string.ascii_uppercase
print(characters)

width, height, n_len, n_class = 90, 30, 4, len(characters)
# width, height, n_len, n_class = 170, 80, 4, len(characters)



from tensorflow.python.keras.models import *
from tensorflow.python.keras.layers import *

# input_tensor = Input((height, width, 3))
# x = input_tensor
# for i, n_cnn in enumerate([2, 2, 2, 2]):
#     for j in range(n_cnn):
#         x = Conv2D(32*2**min(i, 3), kernel_size=3, padding='same', kernel_initializer='he_uniform')(x)
#         x = BatchNormalization()(x)
#         x = Activation('relu')(x)
#     x = MaxPooling2D(2)(x)
#
# x = Flatten()(x)
# x = [Dense(n_class, activation='softmax', name='c%d'%(i+1))(x) for i in range(n_len)]
# model = Model(inputs=input_tensor, outputs=x)
# model.summary()
# exit()

from tensorflow.python.keras.utils import plot_model
from IPython.display import Image

# plot_model(model, to_file='cnn.png', show_shapes=True)
# Image('cnn.png')
#


from tensorflow.python.keras.callbacks import EarlyStopping, CSVLogger, ModelCheckpoint
from tensorflow.python.keras.optimizers import *

def decode(y):
    y = np.argmax(np.array(y), axis=2)[:,0]
    return ''.join([characters[x] for x in y])


if __name__=='__main__':
    # gen_imgs("./img_train/")
    # gen_imgs("./img_test/")
    # valid_data = CaptchaSequence(characters, batch_size=128, steps=5)
    # callbacks = [EarlyStopping(patience=3), CSVLogger('cnn_win.csv'),
    #              ModelCheckpoint('cnn_best_win.h5', save_best_only=True)]
    #
    # model.compile(loss='categorical_crossentropy',
    #               optimizer=Adam(1e-3, amsgrad=True),
    #               metrics=['accuracy'])
    # model.fit_generator(train_data, epochs=100, validation_data=valid_data, workers=4, use_multiprocessing=True,
    #                     callbacks=callbacks)
    #
    # ### 载入最好的模型继续训练一会
    # # 为了让模型充分训练，我们可以载入之前最好的模型权值，然后降低学习率为原来的十分之一，继续训练，这样可以让模型收敛得更好。
    #
    # # model.load_weights('cnn_best_win.h5')
    # # callbacks = [EarlyStopping(patience=3), CSVLogger('cnn_win.csv', append=True),
    # #              ModelCheckpoint('cnn_best_win.h5', save_best_only=True)]
    # #
    # # model.compile(loss='categorical_crossentropy',
    # #               optimizer=Adam(1e-4, amsgrad=True),
    # #               metrics=['accuracy'])
    # # model.fit_generator(train_data, epochs=100, validation_data=valid_data, workers=4, use_multiprocessing=True,
    # #                     callbacks=callbacks)
    model = load_model('cnn_best_szjj.h5', compile=False)
    img_paths_test = [r'D:\vcode3_label']
    data = BaseSequence(paths=img_paths_test, characters=characters, batch_size=1, steps=60, n_len=n_len,
                              width=width, height=height)

    # print(len(data))
    t = 0
    for i in range(len(data)):
        X, y = data[i]
    #     plt.imshow(X[0])
    #     plt.title(decode(y))
    #     plt.show()
    # exit()

        y_pred = model.predict(X)
        real = decode(y)
        pred = decode(y_pred)
        rs = real.__eq__(pred)

        if rs:
            t+=1
        print('real: %s pred:%s rs:%s' % (real, pred, rs))
    print("success count:",str(t))
        # print(y)
        # print(y_pred)
        # plt.title('real: %s\npred:%s' % (decode(y), decode(y_pred)))
        # plt.imshow(X[0], cmap='gray')
        # plt.axis('off')
        # plt.show()